package com.snailframework.gateway.route;

import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.DefaultResponse;
import org.springframework.cloud.client.loadbalancer.EmptyResponse;
import org.springframework.cloud.client.loadbalancer.Request;
import org.springframework.cloud.client.loadbalancer.RequestDataContext;
import org.springframework.cloud.client.loadbalancer.Response;
import org.springframework.cloud.loadbalancer.core.NoopServiceInstanceListSupplier;
import org.springframework.cloud.loadbalancer.core.ReactorServiceInstanceLoadBalancer;
import org.springframework.cloud.loadbalancer.core.SelectedInstanceCallback;
import org.springframework.cloud.loadbalancer.core.ServiceInstanceListSupplier;
import reactor.core.publisher.Mono;

import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;

/**
 * @author: zgadmin(zhangge)
 * @Date: 2023/8/31 15:25
 * @Version: 1.0.0
 */
@Slf4j
public class CustomMetadataLoadBalanceRule implements ReactorServiceInstanceLoadBalancer {

    final String serviceId;

    final AtomicInteger position;

    ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListProvider;

    public CustomMetadataLoadBalanceRule(ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListProvider, String serviceId) {
        this.serviceId = serviceId;
        this.position = new AtomicInteger(new Random().nextInt(1000));
        this.serviceInstanceListProvider = serviceInstanceListProvider;
    }

    @Override
    public Mono<Response<ServiceInstance>> choose(Request request) {
        ServiceInstanceListSupplier supplier = serviceInstanceListProvider.getIfAvailable(NoopServiceInstanceListSupplier::new);

        RequestDataContext context = (RequestDataContext) request.getContext();

        String targetVersion = context.getClientRequest().getHeaders().getFirst("version");

        return supplier.get(request).next().map(serviceInstances -> {
            Response<ServiceInstance> serviceInstanceResponse;
            if (serviceInstances.isEmpty()) {
                if (log.isWarnEnabled()) {
                    log.warn("No servers available for service: {}", serviceId);
                }
                serviceInstanceResponse = new EmptyResponse();
            } else {
                serviceInstanceResponse = versionLoadBalanceServiceInstances(targetVersion, serviceInstances);
            }
            if (supplier instanceof SelectedInstanceCallback && serviceInstanceResponse.hasServer()) {
                ((SelectedInstanceCallback) supplier).selectedServiceInstance(serviceInstanceResponse.getServer());
            }
            return serviceInstanceResponse;
        });
    }

    private Response<ServiceInstance> randomServiceInstances(List<ServiceInstance> serviceInstances) {
        int             pos      = this.position.incrementAndGet() & Integer.MAX_VALUE;
        ServiceInstance instance = serviceInstances.get(pos % serviceInstances.size());
        return new DefaultResponse(instance);
    }

    private Response<ServiceInstance> versionLoadBalanceServiceInstances(String targetVersion, List<ServiceInstance> serviceInstances) {
        List<ServiceInstance> versionServiceInstances = serviceInstances.stream().filter(serviceInstance -> {
            Map<String, String> metadata = serviceInstance.getMetadata();
            if (null == metadata || metadata.isEmpty()) {
                return false;
            }
            String version = metadata.get("version");
            if (null == version || "".equals(version)) {
                return false;
            }
            return version.equals(targetVersion);
        }).collect(Collectors.toList());
        if (null == versionServiceInstances || versionServiceInstances.isEmpty()) {
            return randomServiceInstances(serviceInstances);
        }
        return randomServiceInstances(versionServiceInstances);
    }

}
